{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-06-10T09:58:06.379737Z",
     "start_time": "2024-06-10T09:58:06.108928Z"
    }
   },
   "source": [
    "import os\n",
    "\n",
    "import requests\n",
    "import tiktoken\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import pandas as pd\n",
    "import math\n",
    "\n",
    "if not os.path.exists('sales-textbook.txt'):\n",
    "    url = 'https://huggingface.co/datasets/goendalf666/sales-textbook_for_convincing_and_selling/raw/main/sales_textbook.txt'\n",
    "    with open('sales-textbook.txt', 'wb') as f:\n",
    "        f.write(requests.get(url).content)\n",
    "\n",
    "with open('sales-textbook.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()\n",
    "    # print(text)\n",
    "\n",
    "# 模型的超参数\n",
    "batch_size = 4\n",
    "context_length = 16\n",
    "d_model = 64\n",
    "num_heads = 4\n",
    "\n",
    "# 编码\n",
    "enc = tiktoken.get_encoding(\"o200k_base\")\n",
    "tokenize_text = enc.encode(text)\n",
    "tokenize_text = torch.tensor(tokenize_text, dtype=torch.long)\n",
    "# print(len(tokenize_text))\n",
    "# 最大索引对应的值\n",
    "max_token_value = tokenize_text.max().item()\n",
    "\n",
    "# 分割 训练和验证数据\n",
    "train_text = int(len(tokenize_text) * 0.9)\n",
    "train_data = tokenize_text[:train_text]\n",
    "# print(len(train_data))\n",
    "valid_data = tokenize_text[train_text:]\n",
    "# print(len(valid_data))\n",
    "\n",
    "data = train_data\n",
    "idxs = torch.randint(low=0, high=len(data) - context_length, size=(batch_size,))\n",
    "# print(idxs)\n",
    "x_batch = torch.stack([data[idx:idx + context_length] for idx in idxs])\n",
    "# print(x_batch.shape)\n",
    "y_batch = torch.stack([data[idx + 1:idx + context_length + 1] for idx in idxs])\n",
    "# print(y_batch.shape)\n",
    "\n",
    "df = pd.DataFrame(x_batch[0].numpy())\n",
    "# print(df)\n",
    "# print(enc.decode([28058]))\n",
    "# print(enc.decode(x_batch[0].numpy()))\n",
    "\n",
    "# print(max_token_value)\n",
    "\n",
    "input_embedding_lookup_table = nn.Embedding(max_token_value + 1, d_model)\n",
    "# print(input_embedding_lookup_table)\n",
    "# print(input_embedding_lookup_table.weight.data)\n",
    "\n",
    "x_batch_embedding = input_embedding_lookup_table(x_batch)\n",
    "y_batch_embedding = input_embedding_lookup_table(y_batch)\n",
    "\n",
    "# print(x_batch_embedding.shape)\n",
    "# print(y_batch_embedding.shape)\n",
    "\n",
    "position_encoding_lookup_table = torch.zeros(context_length, d_model)\n",
    "# print(position_encoding_lookup_table)\n",
    "\n",
    "position = torch.arange(0, context_length, dtype=torch.float).unsqueeze(1)\n",
    "# apply the sine & cosine\n",
    "div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
    "position_encoding_lookup_table[:, 0::2] = torch.sin(position * div_term)\n",
    "position_encoding_lookup_table[:, 1::2] = torch.cos(position * div_term)\n",
    "position_encoding_lookup_table = position_encoding_lookup_table.unsqueeze(0).expand(batch_size, -1,\n",
    "                                                                                    -1)  #add batch to the first dimension\n",
    "\n",
    "# print(\"Position Encoding Look-up Table: \", position_encoding_lookup_table.shape)\n",
    "\n",
    "x = x_batch_embedding + position_encoding_lookup_table\n",
    "y = y_batch_embedding + position_encoding_lookup_table\n",
    "# print(pd.DataFrame(x[0].detach().numpy()))\n",
    "# print(pd.DataFrame(y[0].detach().numpy()))\n",
    "\n",
    "# get wq wk wv\n",
    "Wq = nn.Linear(d_model, d_model)\n",
    "Wk = nn.Linear(d_model, d_model)\n",
    "Wv = nn.Linear(d_model, d_model)\n",
    "\n",
    "# get Q K V\n",
    "Q = Wq(x)\n",
    "K = Wk(x)\n",
    "V = Wv(x)\n",
    "\n",
    "# apply mutil head attention\n",
    "Q_ = Q.reshape(batch_size, context_length, num_heads, d_model // num_heads).permute(0, 2, 1, 3)\n",
    "K_ = K.reshape(batch_size, context_length, num_heads, d_model // num_heads).permute(0, 2, 1, 3)\n",
    "V_ = V.reshape(batch_size, context_length, num_heads, d_model // num_heads).permute(0, 2, 1, 3)"
   ],
   "outputs": [],
   "execution_count": 9
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-10T09:54:08.334324Z",
     "start_time": "2024-06-10T09:54:08.327986Z"
    }
   },
   "cell_type": "code",
   "source": "output = Q_ @ K_.transpose(-2, -1) / math.sqrt(d_model // num_heads)",
   "id": "755eb3e3c1a12b67",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-10T09:54:48.802714Z",
     "start_time": "2024-06-10T09:54:48.735009Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# apply mask\n",
    "mask = torch.triu(torch.ones(context_length, context_length), diagonal=1).bool()\n",
    "output = output.masked_fill(mask == 1, float('-inf'))\n",
    "pd.DataFrame(output[0, 0].detach().numpy())"
   ],
   "id": "da0aed33fe352dd4",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "          0         1         2         3         4         5         6   \\\n",
       "0  -0.394613      -inf      -inf      -inf      -inf      -inf      -inf   \n",
       "1  -0.422251  0.384471      -inf      -inf      -inf      -inf      -inf   \n",
       "2  -0.680690 -0.349909 -0.417825      -inf      -inf      -inf      -inf   \n",
       "3  -0.289918  0.357474 -0.090740  0.453864      -inf      -inf      -inf   \n",
       "4  -0.612680  0.294244 -0.673586  0.118284  0.167566      -inf      -inf   \n",
       "5  -0.328054 -0.420368 -1.028461  0.125079  0.024670 -0.042227      -inf   \n",
       "6  -0.371833  0.076859 -0.903844 -0.161742 -0.113930 -0.462978 -0.434631   \n",
       "7  -0.067915 -0.856260 -1.095195 -0.156342 -0.360416 -0.179036  0.707052   \n",
       "8  -0.196987 -0.329827 -0.970322  0.001096  0.573595  0.212441  0.312611   \n",
       "9  -0.128660 -0.202214 -0.515300 -0.168972 -0.023095  0.064054  0.248383   \n",
       "10 -0.463265  0.000148 -0.124091 -0.103220  0.483620  0.167173  0.031562   \n",
       "11 -0.439628 -0.051604  0.071186  0.139874 -0.090942  0.262839 -0.334120   \n",
       "12 -0.018222  0.516332  0.236583  0.538419  0.046674 -0.588419 -0.456478   \n",
       "13 -0.914621  0.687666  0.526141  0.788575 -1.183527 -1.214524 -1.185114   \n",
       "14 -0.802484 -0.006201  0.436912  0.434449  0.114019 -0.107384 -0.343845   \n",
       "15 -0.246566  0.501949  0.195527 -0.059544  0.005731 -0.522259 -0.548661   \n",
       "\n",
       "          7         8         9         10        11        12        13  \\\n",
       "0       -inf      -inf      -inf      -inf      -inf      -inf      -inf   \n",
       "1       -inf      -inf      -inf      -inf      -inf      -inf      -inf   \n",
       "2       -inf      -inf      -inf      -inf      -inf      -inf      -inf   \n",
       "3       -inf      -inf      -inf      -inf      -inf      -inf      -inf   \n",
       "4       -inf      -inf      -inf      -inf      -inf      -inf      -inf   \n",
       "5       -inf      -inf      -inf      -inf      -inf      -inf      -inf   \n",
       "6       -inf      -inf      -inf      -inf      -inf      -inf      -inf   \n",
       "7  -0.121163      -inf      -inf      -inf      -inf      -inf      -inf   \n",
       "8   0.570844  0.673543      -inf      -inf      -inf      -inf      -inf   \n",
       "9  -0.129973  0.370060  0.053102      -inf      -inf      -inf      -inf   \n",
       "10  0.184304  0.005540  0.062094 -0.034956      -inf      -inf      -inf   \n",
       "11 -0.115488  0.015636 -0.335963  0.764183  0.659718      -inf      -inf   \n",
       "12 -0.015946 -0.478854 -0.093232 -0.156782  0.512687  0.404082      -inf   \n",
       "13 -0.718913 -1.084444 -0.935851 -0.820358  0.699424  0.297725 -0.164953   \n",
       "14 -0.232206 -0.704066 -0.762852 -0.581769 -0.143297 -0.283987 -0.319747   \n",
       "15 -0.196677 -0.445704 -0.587220  0.064913  0.177481  0.336449  0.070094   \n",
       "\n",
       "          14        15  \n",
       "0       -inf      -inf  \n",
       "1       -inf      -inf  \n",
       "2       -inf      -inf  \n",
       "3       -inf      -inf  \n",
       "4       -inf      -inf  \n",
       "5       -inf      -inf  \n",
       "6       -inf      -inf  \n",
       "7       -inf      -inf  \n",
       "8       -inf      -inf  \n",
       "9       -inf      -inf  \n",
       "10      -inf      -inf  \n",
       "11      -inf      -inf  \n",
       "12      -inf      -inf  \n",
       "13      -inf      -inf  \n",
       "14 -0.117175      -inf  \n",
       "15  0.104199 -0.059444  "
      ],
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>5</th>\n",
       "      <th>6</th>\n",
       "      <th>7</th>\n",
       "      <th>8</th>\n",
       "      <th>9</th>\n",
       "      <th>10</th>\n",
       "      <th>11</th>\n",
       "      <th>12</th>\n",
       "      <th>13</th>\n",
       "      <th>14</th>\n",
       "      <th>15</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-0.394613</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>-0.422251</td>\n",
       "      <td>0.384471</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>-0.680690</td>\n",
       "      <td>-0.349909</td>\n",
       "      <td>-0.417825</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>-0.289918</td>\n",
       "      <td>0.357474</td>\n",
       "      <td>-0.090740</td>\n",
       "      <td>0.453864</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>-0.612680</td>\n",
       "      <td>0.294244</td>\n",
       "      <td>-0.673586</td>\n",
       "      <td>0.118284</td>\n",
       "      <td>0.167566</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>-0.328054</td>\n",
       "      <td>-0.420368</td>\n",
       "      <td>-1.028461</td>\n",
       "      <td>0.125079</td>\n",
       "      <td>0.024670</td>\n",
       "      <td>-0.042227</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>-0.371833</td>\n",
       "      <td>0.076859</td>\n",
       "      <td>-0.903844</td>\n",
       "      <td>-0.161742</td>\n",
       "      <td>-0.113930</td>\n",
       "      <td>-0.462978</td>\n",
       "      <td>-0.434631</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>-0.067915</td>\n",
       "      <td>-0.856260</td>\n",
       "      <td>-1.095195</td>\n",
       "      <td>-0.156342</td>\n",
       "      <td>-0.360416</td>\n",
       "      <td>-0.179036</td>\n",
       "      <td>0.707052</td>\n",
       "      <td>-0.121163</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>-0.196987</td>\n",
       "      <td>-0.329827</td>\n",
       "      <td>-0.970322</td>\n",
       "      <td>0.001096</td>\n",
       "      <td>0.573595</td>\n",
       "      <td>0.212441</td>\n",
       "      <td>0.312611</td>\n",
       "      <td>0.570844</td>\n",
       "      <td>0.673543</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>-0.128660</td>\n",
       "      <td>-0.202214</td>\n",
       "      <td>-0.515300</td>\n",
       "      <td>-0.168972</td>\n",
       "      <td>-0.023095</td>\n",
       "      <td>0.064054</td>\n",
       "      <td>0.248383</td>\n",
       "      <td>-0.129973</td>\n",
       "      <td>0.370060</td>\n",
       "      <td>0.053102</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>-0.463265</td>\n",
       "      <td>0.000148</td>\n",
       "      <td>-0.124091</td>\n",
       "      <td>-0.103220</td>\n",
       "      <td>0.483620</td>\n",
       "      <td>0.167173</td>\n",
       "      <td>0.031562</td>\n",
       "      <td>0.184304</td>\n",
       "      <td>0.005540</td>\n",
       "      <td>0.062094</td>\n",
       "      <td>-0.034956</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>-0.439628</td>\n",
       "      <td>-0.051604</td>\n",
       "      <td>0.071186</td>\n",
       "      <td>0.139874</td>\n",
       "      <td>-0.090942</td>\n",
       "      <td>0.262839</td>\n",
       "      <td>-0.334120</td>\n",
       "      <td>-0.115488</td>\n",
       "      <td>0.015636</td>\n",
       "      <td>-0.335963</td>\n",
       "      <td>0.764183</td>\n",
       "      <td>0.659718</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>-0.018222</td>\n",
       "      <td>0.516332</td>\n",
       "      <td>0.236583</td>\n",
       "      <td>0.538419</td>\n",
       "      <td>0.046674</td>\n",
       "      <td>-0.588419</td>\n",
       "      <td>-0.456478</td>\n",
       "      <td>-0.015946</td>\n",
       "      <td>-0.478854</td>\n",
       "      <td>-0.093232</td>\n",
       "      <td>-0.156782</td>\n",
       "      <td>0.512687</td>\n",
       "      <td>0.404082</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>-0.914621</td>\n",
       "      <td>0.687666</td>\n",
       "      <td>0.526141</td>\n",
       "      <td>0.788575</td>\n",
       "      <td>-1.183527</td>\n",
       "      <td>-1.214524</td>\n",
       "      <td>-1.185114</td>\n",
       "      <td>-0.718913</td>\n",
       "      <td>-1.084444</td>\n",
       "      <td>-0.935851</td>\n",
       "      <td>-0.820358</td>\n",
       "      <td>0.699424</td>\n",
       "      <td>0.297725</td>\n",
       "      <td>-0.164953</td>\n",
       "      <td>-inf</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>-0.802484</td>\n",
       "      <td>-0.006201</td>\n",
       "      <td>0.436912</td>\n",
       "      <td>0.434449</td>\n",
       "      <td>0.114019</td>\n",
       "      <td>-0.107384</td>\n",
       "      <td>-0.343845</td>\n",
       "      <td>-0.232206</td>\n",
       "      <td>-0.704066</td>\n",
       "      <td>-0.762852</td>\n",
       "      <td>-0.581769</td>\n",
       "      <td>-0.143297</td>\n",
       "      <td>-0.283987</td>\n",
       "      <td>-0.319747</td>\n",
       "      <td>-0.117175</td>\n",
       "      <td>-inf</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>-0.246566</td>\n",
       "      <td>0.501949</td>\n",
       "      <td>0.195527</td>\n",
       "      <td>-0.059544</td>\n",
       "      <td>0.005731</td>\n",
       "      <td>-0.522259</td>\n",
       "      <td>-0.548661</td>\n",
       "      <td>-0.196677</td>\n",
       "      <td>-0.445704</td>\n",
       "      <td>-0.587220</td>\n",
       "      <td>0.064913</td>\n",
       "      <td>0.177481</td>\n",
       "      <td>0.336449</td>\n",
       "      <td>0.070094</td>\n",
       "      <td>0.104199</td>\n",
       "      <td>-0.059444</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-10T09:58:08.876542Z",
     "start_time": "2024-06-10T09:58:08.763067Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# apply softmax\n",
    "attention_score = F.softmax(output, dim=-1)\n",
    "attention_score\n",
    "\n"
   ],
   "id": "da231e5712d9afb",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.3086, 0.6914, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.2708, 0.3770, 0.3522,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0307, 0.1523, 0.1296,  ..., 0.0649, 0.0000, 0.0000],\n",
       "          [0.0350, 0.0776, 0.1208,  ..., 0.0567, 0.0694, 0.0000],\n",
       "          [0.0502, 0.1061, 0.0781,  ..., 0.0689, 0.0713, 0.0605]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.5719, 0.4281, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.4081, 0.3021, 0.2898,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0603, 0.0661, 0.0902,  ..., 0.0765, 0.0000, 0.0000],\n",
       "          [0.0620, 0.0384, 0.0698,  ..., 0.0823, 0.0542, 0.0000],\n",
       "          [0.0293, 0.0481, 0.0483,  ..., 0.0924, 0.0419, 0.0671]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.3260, 0.6740, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.3007, 0.4198, 0.2795,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0552, 0.0381, 0.0687,  ..., 0.1142, 0.0000, 0.0000],\n",
       "          [0.0606, 0.0424, 0.0742,  ..., 0.0684, 0.0440, 0.0000],\n",
       "          [0.0408, 0.0395, 0.0632,  ..., 0.1187, 0.0497, 0.0292]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.5969, 0.4031, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.3848, 0.3401, 0.2751,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.1356, 0.0874, 0.1094,  ..., 0.0488, 0.0000, 0.0000],\n",
       "          [0.1072, 0.0559, 0.0701,  ..., 0.0504, 0.0386, 0.0000],\n",
       "          [0.0707, 0.0777, 0.0518,  ..., 0.0901, 0.0546, 0.0698]]],\n",
       "\n",
       "\n",
       "        [[[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.6375, 0.3625, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.3289, 0.2667, 0.4045,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0395, 0.0568, 0.0604,  ..., 0.0325, 0.0000, 0.0000],\n",
       "          [0.0497, 0.0393, 0.0961,  ..., 0.1026, 0.0488, 0.0000],\n",
       "          [0.0377, 0.0318, 0.1048,  ..., 0.0813, 0.0512, 0.1111]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.4498, 0.5502, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.3641, 0.3114, 0.3245,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0885, 0.0653, 0.0506,  ..., 0.0610, 0.0000, 0.0000],\n",
       "          [0.0355, 0.0570, 0.0965,  ..., 0.0575, 0.0644, 0.0000],\n",
       "          [0.0663, 0.0477, 0.0341,  ..., 0.0733, 0.0684, 0.0743]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.5987, 0.4013, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.2955, 0.3381, 0.3664,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0401, 0.0615, 0.0696,  ..., 0.1748, 0.0000, 0.0000],\n",
       "          [0.0744, 0.0584, 0.0571,  ..., 0.0508, 0.0688, 0.0000],\n",
       "          [0.0288, 0.0834, 0.0435,  ..., 0.1379, 0.0758, 0.0546]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.5059, 0.4941, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.2591, 0.3376, 0.4033,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0811, 0.0609, 0.0606,  ..., 0.0353, 0.0000, 0.0000],\n",
       "          [0.0678, 0.0830, 0.0576,  ..., 0.0336, 0.0555, 0.0000],\n",
       "          [0.0191, 0.0647, 0.0717,  ..., 0.1011, 0.0758, 0.0594]]],\n",
       "\n",
       "\n",
       "        [[[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.3137, 0.6863, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.4609, 0.2255, 0.3137,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0640, 0.0510, 0.0474,  ..., 0.0914, 0.0000, 0.0000],\n",
       "          [0.0547, 0.0914, 0.0415,  ..., 0.0632, 0.0684, 0.0000],\n",
       "          [0.1072, 0.0646, 0.0361,  ..., 0.0794, 0.0685, 0.0454]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.7374, 0.2626, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.2304, 0.4827, 0.2869,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0400, 0.1063, 0.0472,  ..., 0.0488, 0.0000, 0.0000],\n",
       "          [0.0384, 0.0899, 0.0432,  ..., 0.0502, 0.0856, 0.0000],\n",
       "          [0.0696, 0.0587, 0.0419,  ..., 0.0504, 0.0588, 0.0386]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.4410, 0.5590, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.2854, 0.4110, 0.3036,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0636, 0.0609, 0.0887,  ..., 0.0791, 0.0000, 0.0000],\n",
       "          [0.0707, 0.0538, 0.0882,  ..., 0.0625, 0.0515, 0.0000],\n",
       "          [0.0415, 0.0390, 0.0907,  ..., 0.0611, 0.0897, 0.0695]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.5727, 0.4273, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.3457, 0.3185, 0.3357,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0589, 0.0501, 0.0529,  ..., 0.0690, 0.0000, 0.0000],\n",
       "          [0.0196, 0.0419, 0.0825,  ..., 0.0756, 0.0428, 0.0000],\n",
       "          [0.0677, 0.0683, 0.0474,  ..., 0.0417, 0.0297, 0.0524]]],\n",
       "\n",
       "\n",
       "        [[[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.3345, 0.6655, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.3239, 0.4929, 0.1831,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0687, 0.1563, 0.0345,  ..., 0.0526, 0.0000, 0.0000],\n",
       "          [0.1373, 0.1277, 0.0344,  ..., 0.0631, 0.0986, 0.0000],\n",
       "          [0.0515, 0.0965, 0.0464,  ..., 0.0799, 0.0628, 0.0822]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.4982, 0.5018, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.4057, 0.2379, 0.3564,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0639, 0.0505, 0.0707,  ..., 0.0667, 0.0000, 0.0000],\n",
       "          [0.0563, 0.0655, 0.0655,  ..., 0.0761, 0.0633, 0.0000],\n",
       "          [0.0761, 0.0725, 0.0881,  ..., 0.0759, 0.0637, 0.0699]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.6990, 0.3010, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.2829, 0.3924, 0.3247,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0392, 0.0595, 0.0393,  ..., 0.0707, 0.0000, 0.0000],\n",
       "          [0.0669, 0.0663, 0.0606,  ..., 0.0746, 0.0654, 0.0000],\n",
       "          [0.0778, 0.0643, 0.0559,  ..., 0.0356, 0.0753, 0.0513]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.4783, 0.5217, 0.0000,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          [0.4004, 0.3368, 0.2628,  ..., 0.0000, 0.0000, 0.0000],\n",
       "          ...,\n",
       "          [0.0709, 0.0272, 0.0349,  ..., 0.0778, 0.0000, 0.0000],\n",
       "          [0.0953, 0.0912, 0.0451,  ..., 0.0422, 0.0848, 0.0000],\n",
       "          [0.0372, 0.0561, 0.0412,  ..., 0.1101, 0.0803, 0.0406]]]],\n",
       "       grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 10
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-10T10:01:06.742378Z",
     "start_time": "2024-06-10T10:01:06.735596Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# apply attention @ V_\n",
    "A = attention_score @ V_\n",
    "A.shape"
   ],
   "id": "bbbcbd5706d5e42c",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 4, 16, 16])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 13
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-10T10:05:16.036188Z",
     "start_time": "2024-06-10T10:05:16.029206Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# apply concatenate\n",
    "A = A.transpose(1, 2).reshape(batch_size, -1, d_model)\n",
    "Wo = nn.Linear(d_model, d_model)\n",
    "output = Wo(A)\n",
    "output.shape"
   ],
   "id": "8f9b3524a0b39d8c",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 16, 64])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 16
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# apply residual connection\n",
    "output = output + x"
   ],
   "id": "675f107fe0f9de26"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-10T10:13:31.025016Z",
     "start_time": "2024-06-10T10:13:31.020253Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# apply layer normalization\n",
    "layer_norm = nn.LayerNorm(d_model)\n",
    "layer_norm_output = layer_norm(output)\n"
   ],
   "id": "e9d812e0f489b24",
   "outputs": [],
   "execution_count": 21
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-10T10:18:09.123658Z",
     "start_time": "2024-06-10T10:18:09.116537Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# apply feed forward network\n",
    "output = nn.Linear(d_model, d_model * 4)(layer_norm_output)\n",
    "output = nn.ReLU()(output)\n",
    "output = nn.Linear(d_model * 4, d_model)(output)\n",
    "\n",
    "output = output + layer_norm_output\n",
    "output.shape"
   ],
   "id": "236cd925fd44286e",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 16, 64])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 29
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-10T10:18:11.444063Z",
     "start_time": "2024-06-10T10:18:11.440225Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# apply layer normalization\n",
    "output = layer_norm(output)"
   ],
   "id": "c134000551319a48",
   "outputs": [],
   "execution_count": 30
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-10T10:18:20.126773Z",
     "start_time": "2024-06-10T10:18:20.009889Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# apply final linear layer\n",
    "output = nn.Linear(d_model, max_token_value + 1)(output)\n",
    "output.shape"
   ],
   "id": "9c8acc558ba58656",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 16, 199854])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 31
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-06-10T10:23:51.600101Z",
     "start_time": "2024-06-10T10:23:51.532023Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# apply softmax\n",
    "logits = F.softmax(output, dim=-1)\n",
    "pre_index =  torch.argmax(logits[0, 0]).item()\n",
    "enc.decode([pre_index])"
   ],
   "id": "2e0d2b3352bf226b",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'ネ'"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 46
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
